﻿using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Assertions;

public class MonteCarloTreeSearch
{
    public class MonteCarloNode
    {
        public float m_Score = 0;
        public int m_WinCount = 0;
        public int m_LossCount = 0;
        public int m_DrawCount = 0;
        public int m_VisitCount = 0;
        public Game m_Game;
        public bool m_GameOver;
        public Player m_Winner;
        public Player m_Viewer; // 以谁的视角看待胜负
        public Player m_CurrentPlayer;
        public MonteCarloNode m_Parent;
        public Vector2Int m_PlacePos; // 相较父节点落子的位置
        public List<MonteCarloNode> m_Childrens;
        public List<Vector2Int> m_UnExplodedPos;
        public bool FullyExpanded => m_UnExplodedPos.Count == 0;
        public float WinRate => (float)m_WinCount / m_VisitCount;

        public MonteCarloNode()
        {

        }
        public static MonteCarloNode Create(Game game, Player viewer)
        {
            MonteCarloNode node = ObjectPool<MonteCarloNode>.Claim();
            node.m_Game = Game.CopyFrom(game);
            node.m_Viewer = viewer;
            node.m_Score = node.m_WinCount = node.m_LossCount = node.m_DrawCount = node.m_VisitCount = 0;
            node.m_Parent = null;
            node.m_Childrens = ListPool<MonteCarloNode>.Claim();
            node.m_UnExplodedPos = ListPool<Vector2Int>.Claim();
            node.m_UnExplodedPos.AddRange(game.m_Board.m_EmptyCells);
            node.m_CurrentPlayer = game.m_CurrentPlayer;
            node.m_Winner = game.m_Winner;
            node.m_GameOver = game.m_GameOver;
            return node;
        }
        public void Release()
        {
            m_Childrens.Clear();
            ListPool<MonteCarloNode>.Release(m_Childrens);
            m_Parent = null;
            if (m_Game != null)
            {
                m_Game.Release();
                m_Game = null;
            }
            ObjectPool<MonteCarloNode>.Release(this);
            ListPool<Vector2Int>.Release(m_UnExplodedPos);
        }

        /// <summary>
        /// 随机模拟到游戏结束
        /// </summary>
        public void Rollout()
        {
            var tmpGame = Game.CopyFrom(m_Game);
            tmpGame.Rollout();
            m_VisitCount += 1;
            Assert.IsTrue(m_VisitCount == 1);
            if (tmpGame.m_Winner == m_Viewer)
            {
                m_Score += 1;
                m_WinCount += 1;
            }
            else if (tmpGame.m_Winner == Player.None)
            {
                // 平局
                m_DrawCount += 1;
            }
            else
            {
                m_Score -= 1;
                m_LossCount += 1;
            }
            tmpGame.Release();
        }

        /// <summary>
        /// 完全扩展的节点才可以用，用UCB公式找到最佳子节点
        /// </summary>
        /// <returns></returns>
        public MonteCarloNode BestChild()
        {
            float max_ucb = -1;
            int bestActionId = -1;
            int childCount = m_Childrens.Count;
            for (int i = 0; i < childCount; i++)
            {
                var child = m_Childrens[i];
                if (child.m_GameOver)
                {
                    if (child.m_Winner == m_Viewer)
                    {
                        // 如果本身就是必胜的终结态，这个就是最佳的子节点
                        return child;
                    }
                    continue;
                }
                int winCount = m_CurrentPlayer == m_Viewer ? child.m_WinCount : child.m_LossCount; // 最大最小原则
                float ucb = (float)winCount / child.m_VisitCount + Mathf.Sqrt(Mathf.Log((float)m_VisitCount) / child.m_VisitCount);
                if (ucb > max_ucb)
                {
                    max_ucb = ucb;
                    bestActionId = i;
                }
            }
            if (bestActionId < 0 || m_Childrens[bestActionId] == null)
            {
                //for (int j = 0; j < actionCount; j++)
                //{
                //	if (IsActionsValid[j])
                //	{
                //		//Assert.IsTrue(Children[j].Game.actor[0].HP < 1 || Children[j].Game.actor[1].HP < 1 || Game.gameTime >= Game.maxGameTime);
                //		Assert.IsTrue(Children[j].Game.gameover);
                //	}
                //}
                // 结局已成定局
                return null;
                //Debug.Log("???");
            }
            return m_Childrens[bestActionId];
        }

        public void Backpropagate()
        {
            var node = m_Parent;
            while (node != null)
            {
                node.m_Score += m_Score;
                node.m_VisitCount += m_VisitCount;
                node.m_WinCount += m_WinCount;
                node.m_LossCount += m_LossCount;
                node.m_DrawCount += m_DrawCount;
                node = node.m_Parent;
            }
        }
    }

    List<MonteCarloNode> m_nodes;
    public MonteCarloNode m_root;

    public int m_SingleIterationSimulationTimes;
    public int m_iterationTimes;

    public MonteCarloTreeSearch(int singleIterationSimulationTimes = 128)
    {
        m_nodes = new List<MonteCarloNode>(4096);
        m_SingleIterationSimulationTimes = singleIterationSimulationTimes;
    }

    /// <summary>
    /// 根据当前树搜索状态选出一个最佳的落子位置
    /// </summary>
    /// <returns></returns>
    public Vector2Int MakeDecision()
    {
        if (!m_root.FullyExpanded) TreeSearchIteration();

        MonteCarloNode bestChild = null;
        int maxVisitCnt = -1;
        //float maxScore = -9999f;
        foreach (MonteCarloNode child in m_root.m_Childrens)
        {
            Game game = child.m_Game;
            if (child.m_GameOver)
            {
                if (child.m_Winner != m_root.m_Viewer) continue;
                else if (child.m_Winner == m_root.m_Viewer)
                {
                    bestChild = child;
                    break;
                }
            }

            // 蒙特卡洛树搜索的思想：访问最多的一定最好，这是选择节点扩展时就决定的
            if (child.m_VisitCount > maxVisitCnt)
            //if (m_root.Children[i].Score / (float)m_root.Children[i].VisitCount > maxScore)
            {
                bestChild = child;
                maxVisitCnt = child.m_VisitCount;
                //maxScore = m_root.Children[i].Score / (float)m_root.Children[i].VisitCount;
            }
        }
        //if (bestChild != null)
        //    Debug.Log(bestChild.m_VisitCount);
        //else
        //    Debug.Log("bestChild null");
        return bestChild == null ? m_root.m_Childrens[0].m_PlacePos : bestChild.m_PlacePos;
    }

    public void StartNew(Game game, Player viewer)
    {
        Release();
        m_root = MonteCarloNode.Create(game, viewer);
        m_nodes.Add(m_root);
        m_iterationTimes = 0;
    }

    /// <summary>
    /// 选择一个最急切的未被完全扩展的节点
    /// </summary>
    /// <returns></returns>
    MonteCarloNode Select()
    {
        var node = m_root;
        while (node != null && node.m_GameOver == false)
        {
            if (node.m_UnExplodedPos.Count > 0)
            {
                // 当前节点就是未被完全扩展的节点
                return node;
            }
            else
            {
                node = node.BestChild();
            }
        }
        return null;
    }

    /// <summary>
    /// 扩展一个节点
    /// </summary>
    /// <param name="node">扩展的父节点</param>
    /// <returns></returns>
    MonteCarloNode Expand(MonteCarloNode node, Player viewer)
    {
        MonteCarloNode childNode = MonteCarloNode.Create(node.m_Game, viewer);
        childNode.m_Parent = node;
        m_nodes.Add(childNode);
        // take action
        //Assert.IsTrue(node.m_UnExplodedPos.Count > 0);
        int actionId = Random.Range(0, node.m_UnExplodedPos.Count);
        Vector2Int pos = node.m_UnExplodedPos[actionId];
        node.m_UnExplodedPos.RemoveAt(actionId);
        if (node.m_UnExplodedPos.Count == 0)
        {
            // 已经全部扩展，回收游戏资源
            
            node.m_Game.Release();
            node.m_Game = null;
        }
        childNode.m_Game.Step(pos);
        childNode.m_GameOver = childNode.m_Game.m_GameOver;
        childNode.m_Winner = childNode.m_Game.m_Winner;
        childNode.m_PlacePos = pos;
        childNode.m_CurrentPlayer = childNode.m_Game.m_CurrentPlayer;
        node.m_Childrens.Add(childNode);

        return childNode;
    }

    public void TreeSearchIteration()
    {
        for (int i = 0; i < m_SingleIterationSimulationTimes; i++)
        {
            var select_node = Select();
            if (select_node == null)
            {
                // 可能是存在必然获胜的子节点
                //Debug.Log("Select Node Failed");
                m_iterationTimes += 1;
                return;
            }

            //Debug.Log($"select time cost: {DateTime.Now.Second * 1000 + DateTime.Now.Millisecond - startTime}ms");

            var new_node = Expand(select_node, m_root.m_Viewer);
            if (new_node == null)
            {
                //Debug.Log("Expand Node Failed");
                m_iterationTimes += 1;
                return;
            }

            new_node.Rollout();
            new_node.Backpropagate();
        }
        m_iterationTimes += 1;
    }

    // 对象回收给对象池
    public void Release()
    {
        int cnt = m_nodes.Count;
        for (int i = 0; i < cnt; i++)
        {
            ObjectPool<MonteCarloNode>.Release(m_nodes[i]);
        }
        m_nodes.Clear();
    }
}